
from diffusers import StableDiffusionPipeline,StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import torch
from PIL import Image


def main():
    # 加载预训练的 Stable Diffusion 模型
    model_id = "../../../../stable-diffusion-v1-5-text"
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    
    
    # 加载 LoRA 权重
    lora_weights_path = r"D:\lora_train\diffusers-main\diffusers-main\examples\text_to_image\models_japan_cartoon\checkpoint-68500"
    # 使用 PEFT 库的 PeftModel 来加载 LoRA 权重
    pipe.load_lora_weights(lora_weights_path)
    pipe = pipe.to("cuda")
    # 生成图像的提示词
    prompt = "a beautiful chinese girl"
    # 生成图像
    image = pipe(prompt,width=128,height=128).images[0]
    # 保存生成的图像
    image.save("generated_image.jpg")

def main2():
    base_model_path = r"D:\lora_train\stable-diffusion-v1-5-text"
    controlnet_path = r"D:\lora_train\diffusers-main\diffusers-main\examples\controlnet\canny"
    lora_weights_path = r"D:\lora_train\diffusers-main\diffusers-main\examples\text_to_image\models_japan_cartoon\checkpoint-68500"

    controlnet = ControlNetModel.from_pretrained(controlnet_path, torch_dtype=torch.float16)
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        base_model_path, controlnet=controlnet, torch_dtype=torch.float16
    )
    pipe.safety_checker=None
    pipe.load_lora_weights(lora_weights_path)
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
    pipe = pipe.to("cuda")
    control_image = Image.open(r'D:\lora_train\diffusers-main\diffusers-main\examples\controlnet\edges_10001.png').convert("RGB").resize([256,256])
    prompt = "japan cartoon style,a beautiful chinese girl"
    generator = torch.manual_seed(900)
    image = pipe(
        prompt, guidance_scale=7.5,controlnet_conditioning_scale=1.0,num_inference_steps=20, generator=generator, image=control_image,width=256,height=256,
    ).images[0]

    # 创建一个新的空白图像，宽度为两个图像宽度之和，高度不变
    combined_image = Image.new('RGB', (control_image.width + image.width, control_image.height))
    # 粘贴control_image到新图像的左侧
    combined_image.paste(control_image, (0, 0))
    # 粘贴生成的image到新图像的右侧
    combined_image.paste(image, (control_image.width, 0))

    combined_image.save("./output_combined.png")
if __name__ == "__main__":
    main2()